#include "LambertianMaterial.cuh"

using namespace OBR;

/**
 * @brief calc combined effect of bsdf / cos theta
 *
 */
static __forceinline__ __device__ Spectrum calcFr(SamplePRD *prd, float *pdf,
                                                  LambertianMaterialData *data)
{
    InteractionData &si = prd->inter;
    float costheta = max(0.0f, dot(si.wi.d, si.n));
    *pdf = costheta * M_1_PIf;
    float rad = si.has_differentials ? sqrtf(min(si.dudx * si.dudx + si.dudy * si.dudy,
                                                 si.dvdx * si.dvdx + si.dvdy * si.dvdy))
                                     : 0;
    // float uvcolor[3] = {si.uv.x, si.uv.y, 0};
    // return Spectrum(uvcolor) * costheta * M_1_PIf;
    return data->albedo.sample(si.uv, rad, rad, MipmapSampleMode::SQUARE) * costheta * M_1_PIf;
}

/**
 * @brief generate new ray info and add to the next ray, then return beta of the ray
 *
 * @param prd
 * @param instance_id
 * @return beta
 */
static __forceinline__ __device__ Spectrum rayGeneration(SamplePRD *prd, float *pdf,
                                                         LambertianMaterialData *data, float u0,
                                                         float u1)
{
    // sample hemisphere
    Onb onb(prd->inter.n);
    float3 hemisphere_dir = cosine_sample_hemisphere(u0, u1);

    prd->inter.wi = prd->inter.spawnRay(onb.local2world(hemisphere_dir));

#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("==== material raygen\n");
        printf("get hit normal:");
        print_float3(prd->inter.n);
        printf("\n");
        printf("sampled hemisphere:");
        print_float3(hemisphere_dir);
        printf("\n");
        printf("result:");
        print_float3(prd->inter.wi.d);
        printf("\n");
    }
#endif

    return calcFr(prd, pdf, data);
}

extern "C" __device__ Spectrum __direct_callable__lambertian_material(SamplePRD *prd, float *pdf,
                                                                      float u0, float u1,
                                                                      MaterialCallMode mode)
{
    // get sbt data
    CallableData *sbt_record = reinterpret_cast<CallableData *>(optixGetSbtDataPointer());
    LambertianMaterialData *data =
        reinterpret_cast<LambertianMaterialData *>(sbt_record->callable_payload);
    data += prd->inter.mat_instance_id;

    if (mode == MaterialCallMode::SAMPLE_FR)
    {
        return rayGeneration(prd, pdf, data, u0, u1);
    }
    else if (mode == MaterialCallMode::CALC_FR_PDF)
    {
        return calcFr(prd, pdf, data);
    }
    return Spectrum(0.0f);
}